from flask import (Blueprint, request, session)
import psycopg2
import pandas
from flaskr import Utils
import time
import hashlib
from flaskr import Utils
m = hashlib.md5()

bp = Blueprint('postgresSql',__name__, url_prefix='/postgresSql')

@bp.route('', methods = ['GET', 'POST'])
def postgreSql():
    dbname = request.json.get('dbname')
    user = request.json.get('user')
    host = request.json.get('host')
    password = request.json.get('password')
    port = request.json.get('port')

    try:
        con = psycopg2.connect(dbname=dbname, user=user, host=host, password=password, port=port)

    except Exception as e:
        print(e)
        return{
            'data': None,
            'code': 500,
            'msg': '连接数据库异常，请检查信息是否正确'
        }
    data = getDataBase(con, dbname)
    data['database'] = True

    session['data'] = data

    return{
        'data': data,
        'code': 200,
        'msg': '获取数据成功'
    }


def getTableName(con):
    select_table_sql = "select tablename from pg_tables where schemaname='public'"

    tablename_df = pandas.read_sql(select_table_sql, con)

    tablename_list = tablename_df['tablename'].values

    return list(tablename_list)

def getDataBase(con, dbname):
    tablename_list = getTableName(con)
    data_name = dbname
    sheet_table = {}
    sheet_json = {}
    md5s = {}
    for tablename in tablename_list:
        select_table_data_sql = "select * from " + tablename

        generator_ = pandas.read_sql(select_table_data_sql, con,  chunksize=100000)

        df = Utils.generatorToDataFrame(generator_)

        md5 = hashlib.md5((str(time)+tablename).encode(encoding="UTF-8")).hexdigest()
        Utils.addPkl(df, md5)

        md5s[tablename] = md5

        df = df.head(200)

        sheet_table[tablename] = df.to_html()
        sheet_json[tablename] = df.to_json(orient='split')


        # sheets = Utils.getColumns(df)

    data = {
        'dataName': data_name,
        'sheetTable': sheet_table,
        'sheetJson': sheet_json,
        'sheets': tablename_list,
        'md5s': md5s,
        'database': True,
    }

    return data